import threading
import time
import uuid
from typing import Optional

class SystemPromptStore:
    """
    带TTL和最大容量的system_prompt存储，线程安全。
    """
    def __init__(self, max_size=1000, ttl=600):
        self.max_size = max_size
        self.ttl = ttl  # 秒
        self._store = {}  # id: (prompt, timestamp)
        self._lock = threading.Lock()

    def save(self, prompt: str) -> str:
        with self._lock:
            # 清理过期或超量
            self._cleanup()
            prompt_id = str(uuid.uuid4())
            self._store[prompt_id] = (prompt, time.time())
            return prompt_id

    def get(self, prompt_id: str) -> Optional[str]:
        with self._lock:
            item = self._store.get(prompt_id)
            if not item:
                return None
            prompt, ts = item
            if time.time() - ts > self.ttl:
                del self._store[prompt_id]
                return None
            return prompt

    def _cleanup(self):
        # 清理过期
        now = time.time()
        expired = [k for k, (_, ts) in self._store.items() if now - ts > self.ttl]
        for k in expired:
            del self._store[k]
        # 超量清理最早的
        if len(self._store) > self.max_size:
            sorted_items = sorted(self._store.items(), key=lambda x: x[1][1])
            for k, _ in sorted_items[:len(self._store)-self.max_size]:
                del self._store[k]

# 单例
system_prompt_store = SystemPromptStore()